package xjay.ai.emotionalsupport.chatmemory;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.support.GeneratedKeyHolder;
import org.springframework.jdbc.support.KeyHolder;
import org.springframework.stereotype.Repository;

import java.sql.PreparedStatement;
import java.sql.Statement;
import java.util.List;
import java.util.Map;

@Repository
public class JdbcChatMemoryRepository implements ChatMemoryRepository {

    private final JdbcTemplate jdbc; //spring提供的简化数据库操作的工具类
    public JdbcChatMemoryRepository(@Qualifier("mysqlJdbcTemplate") JdbcTemplate jdbcTemplate) {
        this.jdbc = jdbcTemplate; //已经注入了mysql数据库了

        initializeDatabase();

    }

    private void initializeDatabase() {
        //两张表：chat_memory存储对话历史，每次插入两条记录；chat_session存储会话历史，每一次会话都留底
        String createSessionTableSql = """
        CREATE TABLE IF NOT EXISTS chat_session (
            id BIGINT AUTO_INCREMENT PRIMARY KEY,
            session_id BIGINT NOT NULL,
            user_id VARCHAR(50) NOT NULL,
            session_topic VARCHAR(200),
            first_user_message TEXT,
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            UNIQUE KEY unique_session_user (session_id, user_id)
        ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci
        """;

        String createMemoryTableSql = """
        CREATE TABLE IF NOT EXISTS chat_memory (
            id BIGINT AUTO_INCREMENT PRIMARY KEY,
            session_table_id BIGINT NOT NULL,
            message_type VARCHAR(20) NOT NULL,
            content TEXT NOT NULL,
            message_order BIGINT NOT NULL,
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (session_table_id) REFERENCES chat_session(id)
        ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci
        """;

        jdbc.execute(createSessionTableSql);
        jdbc.execute(createMemoryTableSql);
    }
    private Long createNewSession(String userId, String sessionId, String firstMessage) {
        String insertSessionSql = """
        INSERT INTO chat_session (session_id, user_id, first_user_message) 
        VALUES (?, ?, ?)
        """;

        KeyHolder keyHolder = new GeneratedKeyHolder();
        jdbc.update(connection -> {
            PreparedStatement ps = connection.prepareStatement(insertSessionSql, Statement.RETURN_GENERATED_KEYS);
            ps.setLong(1, Long.valueOf(sessionId));
            ps.setString(2, userId);
            ps.setString(3, firstMessage);
            return ps;
        }, keyHolder);
        
        return keyHolder.getKey().longValue();
    }


    public void insertMessage(String userId, String sessionId, String message, Boolean isUserMessage) {

        //检查会话是否存在，如果不存在且是用户消息，则创建新会话
        String checkSessionSql = "SELECT id FROM chat_session WHERE session_id = ? AND user_id = ?";
        List<Long> sessionIds = jdbc.queryForList(checkSessionSql, Long.class, Long.valueOf(sessionId), userId);
        
        Long sessionTableId;
        if (sessionIds.isEmpty()) {
            if (isUserMessage) {
                // 如果是用户消息且会话不存在，创建新会话
                sessionTableId = createNewSession(userId, sessionId, message);
            } else {
                throw new IllegalArgumentException("Cannot insert assistant message without existing session");
            }
        } else {
            sessionTableId = sessionIds.get(0);
        }

        //获取当前会话的最大消息顺序号
        Long maxOrder = getMaxOrder(userId, sessionId);

        //插入新消息到chat_memory
        String insertSql = """
            INSERT INTO chat_memory (session_table_id, message_type, content, message_order) 
            VALUES (?, ?, ?, ?)
            """;
        String messageType = isUserMessage ? "USER" : "ASSISTANT";
        jdbc.update(insertSql, sessionTableId, messageType, message, maxOrder + 1);

    }



    /**
     * note:要根据userId和sessionId拿到最新的对话编号
     * @param userId
     * @param sessionId
     * @return
     */
    public Long getMaxOrder(String userId, String sessionId) {
        String getMaxOrderSql = """
                SELECT COALESCE(MAX(cm.message_order), 0)
                FROM chat_memory cm
                INNER JOIN chat_session cs ON cm.session_table_id = cs.id
                WHERE cs.session_id = ? AND cs.user_id = ?
                """;
        return jdbc.queryForObject(getMaxOrderSql, Long.class, Long.valueOf(sessionId), userId);
    }

    /**
     * note: 根据用户id，会话id，消息序号，拿到一条历史消息文本
     * @param userId
     * @param sessionId
     * @param message_order
     * @return
     */
    public Message getMessage(String userId, String sessionId, int message_order) {
        String getMessageSql = """
                SELECT cm.content, cm.message_type
                FROM chat_memory cm
                INNER JOIN chat_session cs ON cm.session_table_id = cs.id
                WHERE cs.session_id = ? AND cs.user_id = ? AND cm.message_order = ?
                """;
        //String是列名，Object是Long, String，Integer之类的值。一行是一个Map。这里应该只有一行
        Map<String, Object> record =
                jdbc.queryForList(getMessageSql, Long.valueOf(sessionId), userId, Long.valueOf(message_order))
                        .get(0);
        String content = (String) record.get("content");
        String messageType = (String) record.get("message_type");
        Message message;
        if(messageType.equals("ASSISTANT")) {
            message = new AssistantMessage(content);
        } else if(messageType.equals("USER")) {
            message = new UserMessage(content);
        } else {
            throw new IllegalArgumentException("Unknown message type: " + messageType);
        }
        return message;
    }

    @Override
    public List<String> findConversationIds() {
        return List.of();
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        return List.of();
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {

    }

    @Override
    public void deleteByConversationId(String conversationId) {

    }
}
